# coding: utf-8
import pickle
import threading
from random import shuffle

import numpy as np


def copy(target):
    return pickle.loads(pickle.dumps(target))


class MultiToursSolution:
    def __init__(self):
        self.tours = None  # tours
        self.costs = None  # each cost
        self.total = 1e99  # total cost
        self.stags = 0  # stagnation
    
    def __repr__(self):
        return f'tours: {self.tours}, costs: {self.costs}, total: {self.total}'
    
    def init_solution(self, cities, distance, access, depot, m):
        """ This method will generate an initial solution of high quality.
        @param cities: city set
        @param distance: distance matrix
        @param access: accessibility matrix
        @param depot: depot city
        @param m: the number of salesmen
        @return: a solution initiated with random shuffle and optimal insertion
        """
        self.tours = [[depot, depot] for _ in range(m)]
        self.costs = [0 for _ in range(m)]
        V = copy(cities)
        shuffle(V)
        for city in V:
            self.optimal_insert(city, distance, access)
        self.total = sum(self.costs)
        self.stags = 0
        return self
    
    def calc_cost(self, cost_func):
        self.costs = [cost_func(tour) for tour in self.tours]
        self.total = sum(self.costs)
        return self
    
    def optimal_insert(self, city, distance, access, allocated=None, execute=True):
        """
        @param city: [int] the city to be inserted
        @param distance: [float][float] the distance matrix
        @param access: [int][int] the access salesmen group for the city
        @param allocated: [int] allocate a specific salesman to insert the city
        @param execute: [bool] whether performs insertion in place
        @return: solution or optimal insertion cost
        """
        assert self.tours is not None
        salesmen = access(city=city) if allocated is None else [allocated]
        expenses = [[distance[self.tours[x][i - 1]][city]  # add edge a-b
                     + distance[city][self.tours[x][i]]  # add edge b-c
                     - distance[self.tours[x][i - 1]][self.tours[x][i]]  # remove edge a-b
                     for i in range(1, len(self.tours[x]))] for x in salesmen]
        expenses = [(min(x), x.index(min(x))) for x in expenses]
        min_expense = min(expenses, key=lambda x: x[0])
        index = expenses.index(min_expense)
        tour_index, position = salesmen[index], expenses[index][1] + 1
        if execute:
            self.tours[tour_index].insert(position, city)
            self.costs[tour_index] += min_expense[0]
            self.total = sum(self.costs)
            return self
        else:
            return min_expense[0]
    
    def neighborhood(self, copy_prob, distance, cost_func, access):
        new_solution = copy(self)
        new_solution.tours = []
        new_solution.costs = []
        dump_cities = []
        for tour in self.tours:
            copy_cities = np.random.uniform(0, 1, size=len(tour)) > copy_prob
            copy_cities[0], copy_cities[-1] = True, True  # depot must be copied
            new_solution.tours.append(np.array(tour)[copy_cities].tolist())
            dump_cities.extend(np.array(tour)[~copy_cities].tolist())
            new_solution.costs.append(0.0)
        shuffle(dump_cities)
        for city in dump_cities:
            new_solution.optimal_insert(city=city, distance=distance, access=access)
        return new_solution.calc_cost(cost_func=cost_func)
    
    def two_opt(self, distance, depth=1):
        threads = [TwoOptThread(tour, distance, depth) for tour in self.tours]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        self.tours = [t.result for t in threads]
        return self


class BreakPointsSolution:
    def __init__(self):
        self.cities = None  # city sequence
        self.breaks = None  # breakpoints
        self.cost = 1e99  # cost
        self.stags = 0  # stagnation
    
    def __repr__(self):
        return f'cities: {self.cities}, breaks: {self.breaks}, cost: {self.cost}.'
    
    def init_solution(self, cities, distance, access, depot, m):
        """
        @param cities: city set
        @param distance: distance matrix
        @param access: accessibility matrix
        @param depot: depot city
        @param m: the number of salesmen
        @return: a solution initiated with random shuffle and optimal insertion
        """
        self.cities = copy(cities)
        self.cities.remove(depot)
        shuffle(self.cities)
        

class TwoOptThread(threading.Thread):
    def __init__(self, tour, distance, depth=1):
        threading.Thread.__init__(self)
        self.tour = tour
        self.distance = distance
        self.depth = depth
        self.result = None
    
    def run(self) -> None:
        self.result = self.two_opt()
    
    def two_opt(self):
        for _ in range(self.depth):
            for i in range(1, len(self.tour)):
                for j in range(i + 1, len(self.tour)):
                    d1 = self.distance[self.tour[i - 1]][self.tour[i]] + self.distance[self.tour[j - 1]][self.tour[j]]
                    d2 = self.distance[self.tour[i - 1]][self.tour[j - 1]] + self.distance[self.tour[i]][self.tour[j]]
                    if d2 < d1:
                        self.tour = self.tour[:i] + list(reversed(self.tour[i:j])) + self.tour[j:]
        return self.tour


# == a demo to use this class

# from VCMTSP import Instance
# import matplotlib.pyplot as plt
#
# data = Instance(file=f'../../data/TSPLIB_VC/att48.tsp/att48.tsp.txt', m=4)

# solution = MultiToursSolution().init_solution(cities=data.cities,
#                                               distance=data.distance_matrix,
#                                               access=data.access,
#                                               depot=data.depot, m=data.m)
# print(solution)
#
# print(BreakPointsSolution().init_solution(cities=data.cities,
#                                           distance=data.distance_matrix,
#                                           access=data.access,
#                                           depot=data.depot, m=data.m))
